import math
import random

class KMeans:
    def __init__(self, n_clusters=3, max_iter=100, tol=1e-4):
        self.n_clusters = n_clusters  # 簇数量
        self.max_iter = max_iter      # 最大迭代次数
        self.tol = tol                # 收敛阈值
        self.centroids = []           # 质心
        self.labels = []              # 样本标签

    def _distance(self, x1, x2):
        """计算欧氏距离"""
        return math.sqrt(sum((a-b)**2 for a, b in zip(x1, x2)))

    def _init_centroids(self, X):
        """随机初始化质心（从样本中选择）"""
        indices = random.sample(range(len(X)), self.n_clusters)
        self.centroids = [X[i] for i in indices]

    def _assign_clusters(self, X):
        """将样本分配到最近的簇"""
        self.labels = []
        for x in X:
            # 计算到各质心的距离
            distances = [self._distance(x, c) for c in self.centroids]
            # 分配到最近的簇
            self.labels.append(distances.index(min(distances)))

    def _update_centroids(self, X):
        """更新质心（簇内样本均值）"""
        new_centroids = []
        for c in range(self.n_clusters):
            # 取当前簇的所有样本
            cluster_samples = [X[i] for i, label in enumerate(self.labels) if label == c]
            # 计算均值作为新质心
            n_features = len(X[0])
            centroid = [sum(sample[i] for sample in cluster_samples)/len(cluster_samples) 
                       for i in range(n_features)]
            new_centroids.append(centroid)
        
        # 计算质心移动距离
        shift = max(self._distance(old, new) for old, new in zip(self.centroids, new_centroids))
        self.centroids = new_centroids
        return shift

    def fit(self, X):
        """训练模型"""
        self._init_centroids(X)
        
        for _ in range(self.max_iter):
            self._assign_clusters(X)
            shift = self._update_centroids(X)
            if shift < self.tol:  # 质心移动小于阈值则收敛
                break

    def predict(self, X):
        """预测新样本的簇标签"""
        return [self.labels[min(range(len(self.centroids)), 
                   key=lambda i: self._distance(x, self.centroids[i]))] for x in X]

    def wcss(self, X):
        """计算簇内平方和（评估聚类效果）"""
        return sum(self._distance(X[i], self.centroids[label])**2 
                  for i, label in enumerate(self.labels))


# 生成模拟数据
def generate_data(n_samples=300, n_features=2, n_centers=3):
    """生成带聚类结构的模拟数据"""
    random.seed(42)
    data = []
    # 生成每个簇的中心
    centers = [[random.uniform(-10, 10) for _ in range(n_features)] 
              for _ in range(n_centers)]
    
    for _ in range(n_samples):
        # 随机选择一个中心
        center = random.choice(centers)
        # 生成围绕中心的随机点（添加噪声）
        point = [coord + random.gauss(0, 1.5) for coord in center]
        data.append(point)
    return data


# 测试代码
if __name__ == "__main__":
    # 生成2D模拟数据
    X = generate_data(n_samples=300, n_features=2)
    
    # 聚类（3个簇）
    kmeans = KMeans(n_clusters=3)
    kmeans.fit(X)
    
    # 输出结果
    print(f"聚类完成，簇内平方和(WCSS)：{kmeans.wcss(X):.2f}")
    print(f"前5个样本的簇标签：{kmeans.labels[:5]}")
    print(f"质心坐标：{[ [round(c, 2) for c in centroid] for centroid in kmeans.centroids ]}")
